import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import (
    mean_squared_error,
    r2_score,
    mean_absolute_error,
    explained_variance_score,
    max_error,
    mean_absolute_percentage_error,
    accuracy_score,
    f1_score,
    roc_auc_score,
    roc_curve,
    confusion_matrix,
    ConfusionMatrixDisplay
)

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# ---------------------------- 数据预处理 ----------------------------
def load_and_preprocess_data():
    system_data_path = 'D:\\vscode大数据\\.vscode\\长三角\\地点1\\供热历史数据\\地点1_combined.csv'
    indoor_temp_path = 'D:\\vscode大数据\\.vscode\\长三角\\地点1\\室内温度采集数据\\output\\地点1_combined.csv'

    try:
        # 读取系统数据并确保时间列为时间序列类型
        system_data = pd.read_csv(
            system_data_path,
            parse_dates=['时间'],
            dtype={'供温(℃)': 'float32', '回温(℃)': 'float32', '环境温度(℃)': 'float32'}
        )
        # 强制将时间列转换为时间序列类型
        system_data['时间'] = pd.to_datetime(system_data['时间'])
        
        # 读取室内温度数据并确保时间列为时间序列类型
        indoor_temp = pd.read_csv(
            indoor_temp_path,
            parse_dates=['采集时间'],
            dtype={'测点温度(℃)': 'float32'}
        )
        # 强制将时间列转换为时间序列类型
        indoor_temp['采集时间'] = pd.to_datetime(indoor_temp['采集时间'])

        # 列重命名标准化
        system_data.rename(columns={
            '时间': 'timestamp',
            '供温(℃)': 'T_supply',
            '回温(℃)': 'T_return',
            '环境温度(℃)': 'T_out'
        }, inplace=True)

        indoor_temp.rename(columns={
            '采集时间': 'timestamp',
            '测点温度(℃)': 'T_in'
        }, inplace=True)
        # 处理左侧数据中的null值
        system_data.dropna(subset=['timestamp'], inplace=True)
        # 合并数据集
        merged = pd.merge_asof(
            left=system_data.sort_values('timestamp'),
            right=indoor_temp.sort_values('timestamp'),
            on='timestamp'
        )

        # 计算关键特征
        merged['ΔT'] = merged['T_supply'] - merged['T_return']  # 供回水温差
        # 确保时间索引为时间序列类型，以便正确计算时间差
        merged.set_index('timestamp', inplace=True)
        merged['dTdt'] = merged['T_in'].diff() / (merged.index.to_series().diff().dt.total_seconds()/3600)  # 温度变化率
        merged.reset_index(inplace=True)  # 重置索引以便后续操作
        
        # 构造滞后特征
        merged['T_in_lag1'] = merged['T_in'].shift(1)  # T(t)
        
        # 过滤有效数据
        merged = merged[
            (merged['ΔT'] > 0) & 
            (merged['ΔT'] < 20) &  # 合理温差范围
            (merged['dTdt'].abs() < 2)  # 合理温度变化率
        ].dropna()

        return merged.reset_index()
    
    except Exception as e:
        print(f"数据处理错误: {str(e)}")
        raise

# ---------------------------- ARX热力学模型 ---------------------------
class ThermalARXModel:
    def __init__(self):
        """初始化模型参数范围"""
        self.params = None  # [a, b, c, d]
        self.param_bounds = [
            (0.8, 1.0),    # a: 系统惯性系数 (应接近1)
            (-0.5, 0.0),   # b: 室外渗透系数 (通常为负值)
            (0.01, 0.5),   # c: 供回水温差效应系数
            (-2, 2)        # d: 补偿常数
        ]
    
    def predict(self, X):
        """ARX模型预测: T(t+1) = a*T(t) + b*T_out(t) + c*ΔT(t) + d"""
        if self.params is None:
            raise ValueError("模型尚未训练")
        
        a, b, c, d = self.params
        return a * X['T_in_lag1'] + b * X['T_out'] + c * X['ΔT'] + d
    
    def _physics_constraints(self, params):
        """物理约束条件"""
        a, b, c, d = params
        penalty = 0
        
        # 系统惯性约束
        if a < 0.8 or a > 1.0:
            penalty += 1e6 * (a - 0.9)**2
            
        # 室外渗透方向约束
        if b > 0:
            penalty += 1e6
            
        return penalty
    
    def loss_function(self, params, X, y):
        """带物理约束的损失函数"""
        try:
            # 计算预测值
            self.params = params  # 临时设置参数以进行预测
            y_pred = self.predict(X)
            
            # 基础MSE
            mse = mean_squared_error(y, y_pred)
            
            # 物理约束惩罚项
            penalty = self._physics_constraints(params)
            
            # 温度合理性约束
            temp_penalty = 1e4 * np.sum((y_pred < 10) | (y_pred > 30))
            
            # 返回总损失
            return mse + penalty + temp_penalty
        except Exception as e:
            print(f"损失函数错误: {str(e)}")
            return np.inf  # 返回一个很大的值，表示优化应避免这种情况

# ---------------------------- PSO算法 ---------------------------
class PSO:
    def __init__(self, model, X, y, param_bounds, n_particles=30, max_iter=100, w=0.8, c1=2.0, c2=2.0):
        self.model = model
        self.X = X
        self.y = y
        self.param_bounds = param_bounds
        self.n_particles = n_particles
        self.max_iter = max_iter
        self.w = w  # 惯性权重
        self.c1 = c1  # 个体学习因子
        self.c2 = c2  # 社会学习因子
        
        self.dim = len(param_bounds)
        self.X_pso = np.zeros((n_particles, self.dim))
        self.V_pso = np.zeros((n_particles, self.dim))
        self.pbest = np.zeros((n_particles, self.dim))
        self.pbest_fitness = np.full(n_particles, np.inf)
        self.gbest = np.zeros(self.dim)
        self.gbest_fitness = np.inf
        
        # 初始化粒子群
        for i in range(n_particles):
            for d in range(self.dim):
                self.X_pso[i, d] = np.random.uniform(param_bounds[d][0], param_bounds[d][1])
            self.pbest[i] = self.X_pso[i].copy()
            fitness = model.loss_function(self.X_pso[i], X, y)
            self.pbest_fitness[i] = fitness
            if fitness < self.gbest_fitness:
                self.gbest = self.X_pso[i].copy()
                self.gbest_fitness = fitness
    
    def optimize(self):
        for t in range(self.max_iter):
            for i in range(self.n_particles):
                # 更新速度
                r1 = np.random.rand(self.dim)
                r2 = np.random.rand(self.dim)
                self.V_pso[i] = self.w * self.V_pso[i] + self.c1 * r1 * (self.pbest[i] - self.X_pso[i]) + self.c2 * r2 * (self.gbest - self.X_pso[i])
                
                # 更新位置
                self.X_pso[i] += self.V_pso[i]
                
                # 边界检查
                for d in range(self.dim):
                    if self.X_pso[i, d] < self.param_bounds[d][0]:
                        self.X_pso[i, d] = self.param_bounds[d][0]
                    elif self.X_pso[i, d] > self.param_bounds[d][1]:
                        self.X_pso[i, d] = self.param_bounds[d][1]
                
                # 计算适应度
                fitness = self.model.loss_function(self.X_pso[i], self.X, self.y)
                
                # 更新个体最佳
                if fitness < self.pbest_fitness[i]:
                    self.pbest[i] = self.X_pso[i].copy()
                    self.pbest_fitness[i] = fitness
                    # 更新全局最佳
                    if fitness < self.gbest_fitness:
                        self.gbest = self.X_pso[i].copy()
                        self.gbest_fitness = fitness
            
            # 打印迭代信息
            print(f"Iteration {t+1}/{self.max_iter}, Best Fitness: {self.gbest_fitness:.4f}")
        
        # 设置模型参数为全局最佳解
        self.model.params = self.gbest

# ---------------------------- 数据离散化函数 ----------------------------
def discretize_temperature(data, target_col='T_in', threshold=22):
    """将温度值离散化为二分类标签"""
    data[f'{target_col}_binary'] = (data[target_col] >= threshold).astype(int)
    return data

# ---------------------------- 增强评估模块 ---------------------------
def comprehensive_evaluate(model, X, y_reg, y_binary, title):
    """综合评估函数（支持回归和分类指标）"""
    try:
        y_pred_reg = model.predict(X)  # 回归预测值
        y_pred_binary = (y_pred_reg >= 22).astype(int)  # 转化为二分类预测值
    except ValueError as e:
        print(f"预测错误: {str(e)}")
        return None
    
    # 回归指标
    metrics_reg = {
        'MSE': mean_squared_error(y_reg, y_pred_reg),
        'RMSE': np.sqrt(mean_squared_error(y_reg, y_pred_reg)),
        'MAE': mean_absolute_error(y_reg, y_pred_reg),
        'MAPE': mean_absolute_percentage_error(y_reg, y_pred_reg),
        'R²': r2_score(y_reg, y_pred_reg),
        'Explained Variance': explained_variance_score(y_reg, y_pred_reg),
        'Max Error': max_error(y_reg, y_pred_reg)
    }
    
    # 分类指标
    metrics_class = {
        'Accuracy': accuracy_score(y_binary, y_pred_binary),
    }
    
    # 检查是否包含两个类别
    if len(np.unique(y_binary)) == 2:
        metrics_class['AUC'] = roc_auc_score(y_binary, y_pred_reg)
    else:
        print(f"{title}中只有一个类别，无法计算 AUC 分数。")
    
    # 打印指标
    print(f"\n{title}回归指标:")
    for metric_name, metric_value in metrics_reg.items():
        print(f"{metric_name}: {metric_value:.4f}")
    
    print(f"\n{title}分类指标:")
    for metric_name, metric_value in metrics_class.items():
        print(f"{metric_name}: {metric_value:.4f}")
    
    # 时间序列对比
    plt.figure(figsize=(15, 5))
    plt.plot(X['timestamp'], y_reg, label='实际温度', alpha=0.7)
    plt.plot(X['timestamp'], y_pred_reg, label='预测温度', linestyle='--')
    plt.title(f'{title}温度预测对比')
    plt.xlabel('时间')
    plt.ylabel('温度(℃)')
    plt.legend()
    plt.show()
    
    # 残差分析
    residuals = y_reg - y_pred_reg
    plt.figure(figsize=(12, 4))
    plt.scatter(y_pred_reg, residuals, alpha=0.5)
    plt.axhline(0, color='r', linestyle='--')
    plt.title(f'{title}残差分布')
    plt.xlabel('预测温度(℃)')
    plt.ylabel('残差(℃)')
    plt.show()

    # 分类指标可视化
    # 1. 混淆矩阵
    cm = confusion_matrix(y_binary, y_pred_binary)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm)
    disp.plot(cmap=plt.cm.Blues)
    plt.title(f'{title}混淆矩阵')
    plt.show()

    # 2. ROC 曲线
    if len(np.unique(y_binary)) == 2:
        fpr, tpr, _ = roc_curve(y_binary, y_pred_reg)
        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {metrics_class["AUC"]:.2f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title(f'{title} ROC 曲线')
        plt.legend(loc='lower right')
        plt.show()
    else:
        print(f"{title}中只有一个类别，无法绘制 ROC 曲线。")

    return metrics_reg, metrics_class

# ---------------------------- 主流程 ----------------------------
if __name__ == "__main__":
    try:
        # 数据加载
        data = load_and_preprocess_data()
        data = discretize_temperature(data, threshold=22)  # 添加分类标签
        print("有效特征:", data.columns.tolist())
        
        # 准备数据集
        features = ['T_in_lag1', 'T_out', 'ΔT']
        target_reg = 'T_in'  # 回归目标
        target_class = 'T_in_binary'  # 分类目标
        
        # 划分训练测试集
        split_idx = int(len(data)*0.9)
        train = data.iloc[:split_idx]
        test = data.iloc[split_idx:]
        
        # 初始化模型
        model = ThermalARXModel()
        
        # 初始化PSO算法
        pso = PSO(
            model=model,
            X=train[features],
            y=train[target_reg],
            param_bounds=model.param_bounds,
            n_particles=30,
            max_iter=100,
            w=0.8,
            c1=2.0,
            c2=2.0
        )
        
        # 模型训练
        pso.optimize()
        
        # 确保模型已训练
        if model.params is None:
            raise ValueError("模型训练失败，无法进行预测")
        
        # 打印模型公式和参数
        a, b, c, d = model.params
        print("\n=== 模型公式 ===")
        print(f"T(t+1) = {a:.4f} * T(t) + {b:.4f} * T_out(t) + {c:.4f} * ΔT(t) + {d:.2f}")
        
        # 模型评估
        print("\n=== 训练集评估 ===")
        train_metrics_reg, train_metrics_class = comprehensive_evaluate(
            model, 
            train, 
            train[target_reg], 
            train[target_class], 
            "训练集"
        )
        
        print("\n=== 测试集评估 ===")
        test_metrics_reg, test_metrics_class = comprehensive_evaluate(
            model, 
            test, 
            test[target_reg], 
            test[target_class], 
            "测试集"
        )
        
        # 参数物理验证
        print("\n=== 物理合理性验证 ===")
        validations = [
            (0.8 <= a <= 1.0, "系统惯性系数合理"),
            (b <= 0, "室外渗透系数方向正确"),
            (c > 0, "温差效应系数为正"),
            (abs(d) < 3, "补偿常数在合理范围")
        ]
        for condition, desc in validations:
            print(f"{desc}: {'通过' if condition else '不通过'}")

    except Exception as e:
        print(f"运行错误: {str(e)}")